import re

from lib.CountFunc import CountFunc
from lib.CountModel import CountModel
from lib.EquationSolver import EquationSolver
from lib.FunctionManager import FunctionManager  # 自定义函数管理器
from lib.VariableCache import VariableCache

if __name__ == "__main__":
    equation_solver = EquationSolver()
    variable_cache = VariableCache()
    function_manager = FunctionManager()  # 初始化函数管理器

    # 允许的内置函数和常量（与CountModel.evaluate_expression同步）
    allowed_functions = {'abs', 'sin', 'cos', 'tan', 'sqrt', 'exp', 'log', 'gamma'}
    allowed_constants = {'e', 'i', 'pi'}
    # 若CountFunc内存在函数插件,将其添加到allowed_functions中使其与CountModel.evaluate_expression同步
    for func_name, func in CountFunc.__dict__.items():
        if callable(func) and not func_name.startswith('_'):
            allowed_functions.add(func_name)

    while True:
        expression = input("请输入计算公式/方程/函数定义（输入 'q' 退出）: ")
        if expression.lower() == 'q':
            break

        # 处理函数定义（新增正则匹配，支持"f(x)=表达式"格式）
        func_def_pattern = r'^(\w+)\(([\w, ]+)\)\s*=\s*([^:]+)$'  # 匹配不含条件的函数定义，如"f(x)=表达式"或"g(a,b)=a+b"
        func_def_match = re.match(func_def_pattern, expression.strip())
        if func_def_match:
            try:
                func_name, args_str, expr = func_def_match.groups()
                args = [arg.strip() for arg in args_str.split(',') if arg.strip()]
                function_manager.functions[func_name] = (args, expr)  # 直接存储函数定义
                print(f"函数 '{func_name}({args_str})' 定义成功")
            except Exception as e:
                print(f"函数定义失败: {e}")
            continue  # 函数定义处理完成后跳过后续逻辑

        # 原有的def前缀兼容处理（可选保留）
        if expression.strip().startswith('def '):
            func_def = expression[4:].strip()
            try:
                function_manager.parse_function_definition(func_def)
                print(f"函数 '{func_def.split('=')[0].strip()}' 定义成功")
            except Exception as e:
                print(f"函数定义失败: {e}")
            continue

        # 替换表达式中的 ans 为缓存值
        if 'ans' in expression:
            ans_value = variable_cache.get_variable('ans')
            if ans_value is not None:
                ans_str = str(ans_value)
                expression = expression.replace('ans', ans_str)
            else:
                print("缓存中没有 'ans' 变量，请先进行计算。")
                continue

        # 判断是否为方程（含未知变量且无等号时补全）
        match = re.search(r'\b[a-zA-Z]+\b', expression)  # 匹配独立单词（如x、y）
        has_unknown = False
        if match:
            matched_word = match.group().lower()
            # 新增：排除自定义函数名（关键修改）
            has_unknown = (
                    matched_word not in allowed_functions
                    and matched_word not in allowed_constants
                    and matched_word not in function_manager.functions  # 检查是否为自定义函数名
            )
        if has_unknown and '=' not in expression:
            expression = f"{expression}=0"  # 补全为方程


        # 乘号补全（处理数字+字母、右括号+字母的情况）
        def insert_mul(match):
            return f"{match.group(1)}*{match.group(2)}"


        expression = re.sub(r'(-?\d+)([a-zA-Z])', insert_mul, expression)  # 如3x→3*x
        expression = re.sub(r'(\))([a-zA-Z])', insert_mul, expression)  # 如(a+b)x→(a+b)*x

        # 区分方程和计算公式
        if '=' in expression:
            # 解方程
            solutions = equation_solver.solve_equation(expression)
            if solutions:
                print(f"方程的解: {[str(sol) for sol in solutions]}")
                variable_cache.set_variable('ans', solutions[0])
            else:
                print("未找到解。")
        else:
            # 处理自定义函数调用（仅处理用户定义的函数）
            try:
                func_call_pattern = r'^(\w+)\(([^)]+)\)$'  # 严格匹配完整函数调用（如f(3)）
                match = re.match(func_call_pattern, expression)
                if match:
                    func_name, args_str = match.groups()
                    if func_name in function_manager.functions:  # 仅处理用户自定义函数
                        args_values = [CountModel.evaluate_expression(arg.strip()).get_original_number() for arg in
                                       args_str.split(',')]
                        result = function_manager.evaluate_function(func_name, *args_values)
                        print(f"计算结果: {CountModel(result)}")
                        variable_cache.set_variable('ans', CountModel(result))
                        continue
            except Exception as e:
                print(f"函数调用错误: {e}")
                continue

            # 普通计算公式处理（内置函数由CountModel处理）
            result = CountModel.evaluate_expression(expression)
            if result:
                print(f"计算结果: {result}")
                variable_cache.set_variable('ans', result)
